#!usr/bin/env python
# -*- coding: utf-8 -*-

"""给定两个坐标点及成本栅格,计算两点之间的最小成本距离
使用 skimage 的 route_through_array 函数计算最小成本距离, 该函数基于 Dijkstra's 算法

Example:
    $ python cost-distance-skimage.py -s (x, y) -e (x, y) cost.tif out.tif
"""

import sys
import re
from osgeo import gdal, osr
from skimage.graph import route_through_array
import numpy as np

def raster2array(raster_fn):
    """将 GeoTiff 文件以数组的形式读入

    参数:
        raster_fn (str) GeoTiff 文件路径

    返回:
        array 数组

    """
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    array = band.ReadAsArray()

    return array

def coord2pixel_offset(raster_fn, x, y):
    """已知空间坐标点 x, y 求其在 GeoTiff 文件数组中对应的索引

    参数:
        raster_fn (str) GeoTiff 文件路径
        x (float) x 坐标
        y (float) y 坐标

    返回:
        x_offset (int) x 偏移
        y_offset (int) y 偏移
    """
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    origin_x = geotransform[0]
    origin_y = geotransform[3]
    pixel_width = geotransform[1]
    pixel_heigh = geotransform[5]
    x_offset = int((x - origin_x) / pixel_width)
    y_offset = int((y - origin_y) / pixel_heigh)

    return x_offset, y_offset

def create_path(cost_surface_fn, cost_surface_array, start_coord, stop_coord):
    """已知成本栅格个,计算两点之间最短成本距离

    参数:
        cost_surface_fn (str) 成本栅格 GeoTiff 文件路径
        cost_surface_array 成本栅格数组
        start_coord 起点坐标
        stop_coord 终点坐标

    返回:
        path 最短成本路径
        weight 成本权重值
    """
    start_coord_x = start_coord[0]
    start_coord_y = start_coord[1]
    start_index_x, start_index_y = coord2pixel_offset(cost_surface_fn, start_coord_x, start_coord_y)

    stop_coord_x = stop_coord[0]
    stop_coord_y = stop_coord[1]
    stop_index_x, stop_index_y = coord2pixel_offset(cost_surface_fn, stop_coord_x, stop_coord_y)

    indices, weight = route_through_array(cost_surface_array, (start_index_y, start_index_x), (stop_index_y, stop_index_x))
    indices = np.array(indices).T
    path = np.zeros_like(cost_surface_array)
    path[indices[0], indices[1]] = 1

    print(f"weight: {weight}")

    return path, weight

def array2raster(new_raster_fn, raster_fn, array, weight):
    """将数组写入 GeoTiff 文件

    参数:
        new_raster_fn (str) 要保存的最短成本距离 GeoTiff 文件名
        raster_fn (str) 成本栅格 GeoTiff 文件路径
        array 要写入栅格文件的最短成本距离数组
        weight 成本权重值
    """
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    origin_x = geotransform[0]
    origin_y = geotransform[3]
    pixel_width = geotransform[1]
    pixel_heigh = geotransform[5]
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    out_raster = driver.Create(new_raster_fn, cols, rows, 1, gdal.GDT_Byte)
    out_raster.SetGeoTransform((origin_x, pixel_width, 0, origin_y, 0, pixel_heigh))
    out_band = out_raster.GetRasterBand(1)
    out_band.SetMetadata({'WEIGHT': str(weight * pixel_width)})
    out_band.WriteArray(array)
    out_raster_srs = osr.SpatialReference()
    out_raster_srs.ImportFromWkt(raster.GetProjectionRef())
    out_raster.SetProjection(out_raster_srs.ExportToWkt())
    out_band.FlushCache()

def main(cost_surface_fn, output_path_fn, start_coord, stop_coord):
    cost_surface_array = raster2array(cost_surface_fn)
    path_array, weight = create_path(cost_surface_fn, cost_surface_array, start_coord, stop_coord)
    array2raster(output_path_fn, cost_surface_fn, path_array, weight)

if __name__ == "__main__":
    argv = sys.argv

    cost_surface_fn = None
    output_path_fn = None
    start_coord = None
    stop_coord = None

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-s' or arg == '-start':
            i = i + 1
            start_coord = re.findall("\d+\.\d+", argv[i])
            start_coord = tuple([float(i) for i in start_coord])

        elif arg == '-e' or arg == '-end':
            i = i + 1
            stop_coord = re.findall("\d+\.\d+", argv[i])
            stop_coord = tuple([float(i) for i in stop_coord])

        elif cost_surface_fn is None:
            cost_surface_fn = argv[i]

        elif output_path_fn is None:
            output_path_fn = argv[i]

        else:
            print("python script.py -s '(x, x)' -e '(x, x)' cost.tif out.tif")
            sys.exit(1)

        i = i + 1

    # cost_surface_fn = 'cost.tif'
    # start_coord = (112.2262, 37.8275)
    # stop_coord = (112.3015, 37.7514)
    # output_path_fn = 'OUT.tif'
    if output_path_fn is None:
        output_path_fn = 'OUT.tif'

    print(f"start from: {start_coord}, end to: {stop_coord}")

    main(cost_surface_fn, output_path_fn, start_coord, stop_coord)